misc: Update gemm/batched gemm cubins from trtllm-gen, gemm header refactor#2740
misc: Update gemm/batched gemm cubins from trtllm-gen, gemm header refactor#2740aleozlx merged 12 commits intoflashinfer-ai:mainfrom
Conversation
jimmyzho
requested review from
IwakuraRein,
bkryu,
cyx-6,
jiahanc,
kahyunnam,
nv-yunzheq and
yzh119
as code owners
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughUpdated TRT‑LLM GEMM integration: kernel name strings adjusted, GEMM problem-dimension validity fields enabled, many exported GEMM headers/types removed, JIT cubin/header download re-enabled with updated artifact paths/checksums, and small test tolerance tweaks. Changes
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request focuses on updating and refining the integration of TRTLLM-Gen GEMM cubins, particularly enhancing support for FP8 operations with a new DeepSeek option. It streamlines kernel selection, enables previously disabled problem dimension validations, and addresses JIT compilation issues to ensure the use of the latest and most optimized kernels. Minor adjustments to test tolerances were also made. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the GEMM cubins from trtllm-gen. The changes involve updating artifact paths and checksums, re-enabling header downloads, and adjusting the C++ code to accommodate the new cubins. A key change is the removal of a hardcoded kernel selection function in favor of a more general tactic selection mechanism. Several TODOs related to trtllm-gen fixes have been addressed, and the corresponding code has been enabled. Test tolerances were also slightly adjusted, which is expected with new kernels. The changes are consistent and look correct. I have a few minor suggestions for code cleanup.
| # "b55211623be7f5697c5262ffd8361fc06c147bc9/batched_gemm-b3c1646-c111d7c/" | ||
| "c6de39dedecf3f6a06072082a23c959c9bdcbf11/batched_gemm-4daf11e-c111d7c/" | ||
| ) | ||
| TRTLLM_GEN_GEMM: str = ( | ||
| "1fddc48b7b48af33914d040051b3e2ee9ba4701e/gemm-145d1b1-9b113e3/" | ||
| # "1fddc48b7b48af33914d040051b3e2ee9ba4701e/gemm-145d1b1-9b113e3/" | ||
| "5b85f1064103f8d01b5aa574b45a355ccee0a635/gemm-4daf11e-8adb015/" |
There was a problem hiding this comment.
The old artifact paths for TRTLLM_GEN_BMM and TRTLLM_GEN_GEMM are commented out. It would be cleaner to remove this commented code before this pull request is merged.
| # "0af823880730c4f0b3832d2208fab035946694b83444410b9309db5613d60195" | ||
| "3c4070eb507f7a8391051ada620756259eae4551dd0a0d797d018d6ae5b7d39f" | ||
| ) | ||
| DEEPGEMM: str = "1a2a166839042dbd2a57f48051c82cd1ad032815927c753db269a4ed10d0ffbf" | ||
| TRTLLM_GEN_GEMM: str = ( | ||
| "15cb8c85dfb5eddd4f121d64cb5a718321fb55b85aa19df10ddc1329d4a726b9" | ||
| # "15cb8c85dfb5eddd4f121d64cb5a718321fb55b85aa19df10ddc1329d4a726b9" | ||
| "6ec458b818f2de27176950598292ed64816397988031d5db0c6799dffc378cf9" |
There was a problem hiding this comment.
The old checksums for TRTLLM_GEN_BMM and TRTLLM_GEN_GEMM are commented out. It would be cleaner to remove this commented code before this pull request is merged.
|
/bot run |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
flashinfer/jit/gemm/core.py (1)
544-554: Consider extracting common header download logic.The header download logic in
gen_trtllm_gen_gemm_module(lines 544-554) andgen_trtllm_low_latency_gemm_module(lines 710-720) is identical. Consider extracting this into a helper function to reduce duplication.♻️ Suggested helper function
def _download_gemm_headers(checksum: bytes) -> None: """Download TRTLLM GEMM headers to the cubin directory.""" include_path = f"{ArtifactPath.TRTLLM_GEN_GEMM}/include" header_path = f"{include_path}/trtllmGen_gemm_export" header_dest_dir = ( jit_env.FLASHINFER_CUBIN_DIR / "flashinfer" / "trtllm" / "gemm" / "trtllmGen_gemm_export" ) download_trtllm_headers( "gemm", header_dest_dir, header_path, ArtifactPath.TRTLLM_GEN_GEMM, checksum )Also applies to: 710-720
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/jit/gemm/core.py` around lines 544 - 554, Extract the duplicated header download block from gen_trtllm_gen_gemm_module and gen_trtllm_low_latency_gemm_module into a small helper (e.g. _download_gemm_headers(checksum: bytes)) that builds include_path/header_path/header_dest_dir and calls download_trtllm_headers("gemm", header_dest_dir, header_path, ArtifactPath.TRTLLM_GEN_GEMM, checksum); then replace the duplicated blocks in gen_trtllm_gen_gemm_module and gen_trtllm_low_latency_gemm_module with a call to that helper to remove duplication.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@csrc/trtllm_low_latency_gemm_runner.cu`:
- Line 172: The code references configOptions.mUseShuffledMatrix which does not
exist; replace that reference with the correct field name
configOptions.mUseShuffledMatrixA (e.g., in the conditional around
configOptions.mUseShuffledMatrix) so it matches the GemmOptions.h definition;
search for other uses of mUseShuffledMatrix in trtllm_low_latency_gemm_runner.cu
and update them to mUseShuffledMatrixA to avoid compilation errors.
---
Nitpick comments:
In `@flashinfer/jit/gemm/core.py`:
- Around line 544-554: Extract the duplicated header download block from
gen_trtllm_gen_gemm_module and gen_trtllm_low_latency_gemm_module into a small
helper (e.g. _download_gemm_headers(checksum: bytes)) that builds
include_path/header_path/header_dest_dir and calls
download_trtllm_headers("gemm", header_dest_dir, header_path,
ArtifactPath.TRTLLM_GEN_GEMM, checksum); then replace the duplicated blocks in
gen_trtllm_gen_gemm_module and gen_trtllm_low_latency_gemm_module with a call to
that helper to remove duplication.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e16a0bf2-3331-4355-bf55-e02ac48a0663
📒 Files selected for processing (6)
csrc/trtllm_gemm_runner.cucsrc/trtllm_low_latency_gemm_runner.cuflashinfer/artifacts.pyflashinfer/jit/cubin_loader.pyflashinfer/jit/gemm/core.pytests/moe/test_trtllm_gen_fused_moe.py
| configOptions.mTransposeMmaOutput == true && | ||
| configOptions.mLayoutA == gemm::gemm::MatrixLayout::BlockMajorK && | ||
| configOptions.mUseShuffledMatrixA) { | ||
| configOptions.mUseShuffledMatrix) { |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for the actual field name in the codebase
rg -n "mUseShuffledMatrix" --type-add 'header:*.h' -t header -C2Repository: flashinfer-ai/flashinfer
Length of output: 4636
🏁 Script executed:
rg -n "mUseShuffledMatrix" csrc/trtllm_low_latency_gemm_runner.cu -C 3Repository: flashinfer-ai/flashinfer
Length of output: 389
Fix field name: mUseShuffledMatrix should be mUseShuffledMatrixA.
At line 172, the code references configOptions.mUseShuffledMatrix, but the GemmOptions.h header defines this field as mUseShuffledMatrixA. Using the incorrect field name will cause a compilation error.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@csrc/trtllm_low_latency_gemm_runner.cu` at line 172, The code references
configOptions.mUseShuffledMatrix which does not exist; replace that reference
with the correct field name configOptions.mUseShuffledMatrixA (e.g., in the
conditional around configOptions.mUseShuffledMatrix) so it matches the
GemmOptions.h definition; search for other uses of mUseShuffledMatrix in
trtllm_low_latency_gemm_runner.cu and update them to mUseShuffledMatrixA to
avoid compilation errors.
|
[FAILED] Pipeline #45774302: 1/20 passed |
jimmyzho
requested review from
aleozlx,
saltyminty,
samuellees,
sricketts,
yongwww and
yyihuang
as code owners
jimmyzho
changed the title
|
/bot run |
|
/bot run |
|
/bot run |
|
/bot run |
|
[SUCCESS] Pipeline #47683475: 10/20 passed |
5 participants
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Bug Fixes
Behavioral Changes
Chores
Tests
Removed